# extract_trained_states.py

import json
from data_loader import load_evaluation_cache, load_agent_paths

# 导入超参数映射
from hyperparameter_mapping import learning_rate_mapping, max_depth_mapping, n_estimators_mapping

def extract_trained_states(cache, paths):
    """
    提取训练过程中访问过的状态及其对应的值。

    参数:
    - cache (dict): 超参数组合到准确率的映射。
    - paths (list): 包含每个Episode路径的列表。

    返回:
    - trained_states (list of dict): 每个元素包含超参数组合和对应的准确率。
    """
    trained_states = []
    seen_states = set()

    for episode_num, path in enumerate(paths, start=1):
        for state in path:
            # 将索引映射为实际的超参数值
            lr_index, md_index, ne_index = state
            lr = learning_rate_mapping.get(lr_index, None)
            md = max_depth_mapping.get(md_index, None)
            ne = n_estimators_mapping.get(ne_index, None)

            if lr is None or md is None or ne is None:
                print(f"Warning: Invalid hyperparameter indices in state {state} of Episode {episode_num}.")
                continue

            state_tuple = (lr, md, ne)
            if state_tuple not in seen_states:
                seen_states.add(state_tuple)
                state_dict = {
                    'learning_rate': lr,
                    'max_depth': md,
                    'n_estimators': ne
                }
                # 将状态字典转换为JSON字符串，确保键排序一致
                state_str = json.dumps(state_dict, sort_keys=True)
                accuracy = cache.get(state_str, None)
                if accuracy is not None:
                    state_dict['accuracy'] = accuracy
                    trained_states.append(state_dict)
                else:
                    print(f"Warning: State {state_dict} not found in evaluation_cache.")
    print(f"Extracted {len(trained_states)} unique trained states.")
    return trained_states

if __name__ == "__main__":
    cache = load_evaluation_cache('evaluation_cache.json')
    paths = load_agent_paths('agent_paths.json')
    trained_states = extract_trained_states(cache, paths)

    # 保存提取的数据以便后续使用
    with open('trained_states.json', 'w') as f:
        json.dump(trained_states, f, indent=4)
    print("Saved trained states to trained_states.json.")
